import cv2
import numpy as np
import xml.etree.ElementTree as ET
from PIL import Image


def label_vis(img_cop, cls, bbox):
    cv2.rectangle(img_cop, (bbox[0], bbox[2]), (bbox[1], bbox[3]), (0, 0, 255), thickness=3)
    cv2.putText(img_cop,cls,(bbox[0], bbox[2]-5),cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 0, 255), 2)
    return img_cop


def read_label(image_id):

    label_file = open('sandbox/Annotations/%s.xml' % image_id)
    img_path = "sandbox/image/%s.jpg" % image_id
    # img = np.array(cv2.imread(img_path))
    img = cv2.imread(img_path)
    img_cop = img
    tree = ET.parse(label_file)
    root = tree.getroot()
    for obj in root.iter('object'):
        cls = obj.find('name').text
        xmlbox = obj.find('bndbox')
        bbox = [int(xmlbox.find('xmin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymin').text),
             int(xmlbox.find('ymax').text)]
        img_cop = label_vis(img_cop, cls, bbox)
    cv2.imwrite("sandbox/image_vis/%s.jpg" % image_id, img_cop)


def main():
    image_ids=open('sandbox/train1.txt').read().strip().split()
    for image_id in image_ids:
        read_label(image_id)


if __name__ == "__main__":
    main()


